import os
import json
import argparse

from v2.dataset_metadata import DATA_FILES


def extract_data_range(input_dir, output_dir, start_line, end_line):
    """
    提取数据集中指定范围的数据

    参数:
    input_dir: 输入数据集目录
    output_dir: 输出目录
    start_line: 起始行号(包含，从0开始)
    end_line: 终止行号(不包含)
    """

    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    for filename in DATA_FILES:
        input_path = os.path.join(input_dir, filename)
        output_path = os.path.join(output_dir, filename)

        # 检查输入文件是否存在
        if not os.path.exists(input_path):
            print(f"警告: 文件 {input_path} 不存在，跳过")
            continue

        print(f"处理文件: {filename}")

        try:
            # 根据文件类型选择处理方式
            if filename.endswith('.txt'):
                if 'face_attrs' in filename or 'audio_features' in filename or 'title' in filename or 'video_features' in filename:
                    # JSON格式文件
                    process_json_file(input_path, output_path, start_line, end_line)
                else:
                    # 普通文本文件
                    process_text_file(input_path, output_path, start_line, end_line)

            print(f"成功提取 {filename} 的第 {start_line}-{end_line - 1} 行")

        except Exception as e:
            print(f"处理文件 {filename} 时出错: {str(e)}")


def process_text_file(input_path, output_path, start_line, end_line):
    """处理普通文本文件"""
    with open(input_path, 'r', encoding='utf-8') as f_in:
        lines = f_in.readlines()

    # 提取指定范围的行
    selected_lines = lines[start_line:end_line]

    with open(output_path, 'w', encoding='utf-8') as f_out:
        f_out.writelines(selected_lines)


def process_json_file(input_path, output_path, start_line, end_line):
    """处理JSON格式文件"""
    data = []
    with open(input_path, 'r', encoding='utf-8') as f_in:
        for line_num, line in enumerate(f_in):
            if start_line <= line_num < end_line:
                try:
                    json_data = json.loads(line.strip())
                    data.append(json_data)
                except json.JSONDecodeError as e:
                    print(f"警告: 第 {line_num} 行JSON解析错误: {str(e)}")

    with open(output_path, 'w', encoding='utf-8') as f_out:
        for item in data:
            f_out.write(json.dumps(item, ensure_ascii=False) + '\n')


def main():
    parser = argparse.ArgumentParser(description='提取TikTok数据集中指定范围的数据')
    parser.add_argument('--input_dir', type=str, default='../resources/dataset/origin',
                        help='输入数据集目录路径，默认: ../resources/dataset/origin')
    parser.add_argument('--output_dir', type=str, default='../resources/dataset/top1000',
                        help='输出目录路径，默认: ../resources/dataset/top1000')
    parser.add_argument('--start_line', type=int, default=0, help='起始行号(从0开始)，默认: 0')
    parser.add_argument('--end_line', type=int, default=1000, help='终止行号(不包含)，默认: 1000')
    args = parser.parse_args()

    args.output_dir = "../resources/dataset/top1000000"
    args.end_line = 1000000

    # 参数验证
    if args.start_line < 0:
        raise ValueError("起始行号必须大于等于0")
    if args.end_line <= args.start_line:
        raise ValueError("终止行号必须大于起始行号")

    # 执行数据提取
    extract_data_range(args.input_dir, args.output_dir, args.start_line, args.end_line)
    print(f"数据提取完成！输出目录: {args.output_dir}")


if __name__ == "__main__":
    main()
